from cmath import log
from itertools import count
from sys import prefix
from tkinter import dialog
import requests
import json
import time
import os
import random


def get_result(input):
    url = 'http://127.0.0.1:8080/api/v0/generate'
    req_data = input
    rsp = requests.post(url, json=req_data)
    if rsp.status_code ==200:
        rsp_data = rsp.json()
        # print(rsp_data)
    else:
        print(rsp.status_code)
    return rsp_data.get('texts')[0]

SPEAKER1 = ' [Human]:'#speaker 前要加一个空格不然，tokenizer会识别成两个字
SPEAKER2 = ' [Robot]:'
Intruction = "This is a discussion between a [Human] and a [Robot]. The [Robot] is very kind, empathtic and humous." + "\n"*5
#1# Sample = SPEAKER1 + "Otters are cute, this no one can deny. They have big eyes and snub snouts and paws like tiny leedle hands.\n"+\
#          SPEAKER2 + "ANIMALS\n"+\
#          SPEAKER1 + "Free, universal preschool for three- and four-year-olds is a key component of the Democrats’ agenda.\n"+\
#          SPEAKER2 + "EDUCATION\n"+\
#          SPEAKER1 + "In normal circumstances the Intergovernmental Panel on Climate Change (IPCC) can take media attention for granted.\n"+\
#          SPEAKER2 + "ENVIRONMENT\n"+\
#          SPEAKER1 + "Burberry has become the latest luxury brand to temporarily shut its stores in Russia following Moscow’s in\n"+\
#          SPEAKER2 + "FASHION\n"+\
#          SPEAKER1 + "Alibaba ‘s quarterly profit is expected to fall almost 60% year over year when the Chinese tech giant reports earnings on Thursday.\n"+\
#          SPEAKER2 + "BUSINESS\n"+\
#          SPEAKER1 + "Most EVS operate at 400 volts (400v). But a number of producers and their component suppliers are now gearing up to introduce 800v drive systems.\n"+\
#          SPEAKER2 + "TECH\n"+\
#          SPEAKER1 + "With only three weeks remaining in the 2021-22 NBA regular season, teams are starting to get down to the final 10-or-so games on their schedule.\n"+\
#          SPEAKER2 + "SPORT\n"
Sample = SPEAKER1 + "How do we learn?\n"+SPEAKER2 + "Through examining our mistakes.\n\n"+SPEAKER1 + "How do we get into the flow?\n"+SPEAKER2 + "By letting go of our ego and self absorption.\n\n"+SPEAKER1 + "How do we grow?\n"+SPEAKER2 + "Through connection and diversity.\n\n"+SPEAKER1 + "What is the purpose of suffering?\n"+SPEAKER2 + "To learn and grow.\n\n"
prompt = Intruction + Sample

max_number_turns = 10
config = {"info":{},"dialogue":[],"meta":[],"user":[],"assistant":[],"user_memory":[]}
info = {"prompt":prompt,
            "max_length":30,
            'do_sample':True,
            "top_p":0.7,
            "top_k":10,
            "temperature":0.9,
            "stop_words":['\n','<|endoftext|>','['],
            "min_length":5,
            "length_penalty":1.0,
            "repetition_penalty":1.3,
            "num_return_sequences":2
}
# info = dict(
#     prompt=prompt,
#     min_length=5,
#     max_length=30,
#     temperature=0.9,
#     do_sample=True,
#     top_p=0.7,
#     top_k=10,
#     stop_words = ['\n','<|endoftext|>','['],
#     length_penalty=1.0,
#     repetition_penalty=1.3,
#     num_return_sequences=2
# )

config["info"] = info
dialogue = []
meta = []
with_knowledge = 3 #先验知识
user_utt = ''
initial_situation = config["info"]["prompt"]
topic_list = ["navigate","schedule","weather"]
emotion_list = ["afraid","angry","annoyed","anticipating","anxious","apprehensive","ashamed","caring","confident","content","devastated","disappointed","disgusted","embarrassed","excited","faithful","furious","grateful","guilty","hopeful","impressed","jealous","joyful","lonely","nostalgic","prepared","proud","sad","sentimental","surprised","terrified","trusting"]
#emotion数据集中meta I是用户

#1
# while user_utt != 'quit':
#     dialog_pairs = []
#     user_utt = input(SPEAKER1)
#     if (user_utt=='quit' or user_utt==''):#创建数据json文件
#         config['dialogue'] = dialogue
#         localtime = time.localtime(time.time())
#         if not os.path.exists(os.path.join("./","data")):
#             os.makedirs(os.path.join("./","data"))
#         with open('./data/%s-%s-%s:%sdata.json'%(localtime.tm_mon,localtime.tm_mday,localtime.tm_hour,localtime.tm_min),'w') as f:
#             json.dump(config, f, indent='\t')
#         break
#     prompt = prompt + SPEAKER1 + user_utt + '\n' + SPEAKER2 # prompt = prompt + Human+'\n'
#     print('='*72)
#     print(prompt)
#     print('='*72)
#     dialog_pairs.append(SPEAKER1 + user_utt)
#     info['prompt'] = prompt
#     answer = get_result(info)[1:-1]
#     dialog_pairs.append(SPEAKER2 + answer.strip('\\n'))
#     print(SPEAKER2 + answer.strip('\\n'))
#     prompt=prompt+answer.strip('\\n')+"\n"
#     dialogue.append(dialog_pairs)

#2
history_dialog = "History Dialogue:\n"
while user_utt != 'quit': 
    dialog_pairs = []
    user_utt = input(SPEAKER1)
    if (user_utt=='quit' or user_utt==''):
        localtime = time.localtime(time.time())
        if not os.path.exists(os.path.join("./","data")):
             os.makedirs(os.path.join("./","data"))
        with open('./data/%s-%s-%s:%sdata.json'%(localtime.tm_mon,localtime.tm_mday,localtime.tm_hour,localtime.tm_min),'w') as f:
             json.dump(config, f, indent='\t')
        break
    
    # dialogue["dialogue"].append([user_utt,""]) #append整个[user_utt,""]
    for turn in config["dialogue"]:
        history_dialog += f"{turn[0]}" +"\n"
        if turn[1] == "":
            ## NO GENERATION REQUIRED
            pass
        else:
            history_dialog += f"{turn[1]}" +"\n\n"
    # print("history:"+'\n'+history_dialog)
    prefix = initial_situation + history_dialog + SPEAKER1 + user_utt + '\n' + SPEAKER2 # prompt = prompt + Human+'\n'
    print('='*72)
    print(prefix)
    print('='*72)
    dialog_pairs.append(SPEAKER1 + user_utt)
    info['prompt'] = prefix

    # response = json.loads(get_result(info))
    sample_times = 5
    for i in range(sample_times):
        response = get_result(info)
        print(f"{i+1}:{response}")
    
    # for swn,sw in enumerate(config["info"]["stop_words"]):
    #     if(response[-1]==sw):
    #         sws = 1
    #         break
    #     sws = 0...

    #history_dialog = history_dialog + SPEAKER1 + user_utt + '\n' + SPEAKER2 + response.strip('\\n') +'\n'
    dialog_pairs.append(SPEAKER2 + response.strip('\\n'))
    print(SPEAKER2 + response.strip('\\n') + "\n")
    dialogue.append(dialog_pairs)


    for topicnum,topic in enumerate(topic_list):
        if topic in user_utt:
            f = open(f"openai/test/data/smd/{topic}-valid.json","r")
            data = json.load(f)
            random.Random(1).shuffle(data)
            sample = data[0]
            prefix = "KB:\n"
            for s in sample["meta"]:
                prefix += s+"\n"
        
            for dialogue_shot in sample["dialogue"][:with_knowledge]:
                prefix += "Dialogue:\n"
                for turn in dialogue_shot:
                    dialog_pairs = []
                    dialog_pairs.append(SPEAKER1 + f"{turn[0]}")
                    prefix += SPEAKER1 + f"{turn[0]}" +"\n"
                    if turn[1] == "":
                        dialog_pairs.append(SPEAKER2 + f"")
                        prefix += SPEAKER2 + f"" 
                    else:
                        dialog_pairs.append(SPEAKER2 + f"{turn[1]}")
                        prefix += SPEAKER2 + f"{turn[1]}" +"\n"
                prefix += "\n\n"

            info['prompt'] = prefix
            dialogue.append(dialog_pairs)
    
    for emotionnum,emotion in enumerate(emotion_list):
        if emotion in user_utt:
            f = open(f"openai/test/data/EmpatheticDialogues/{emotion}.json","r")
            data = json.load(f)
            random.Random(1).shuffle(data)
            introduction = SPEAKER1 +"\n"
            prefix = "Dialogue:\n"
            for sample in data[:with_knowledge]:
                # prefix = "KB:\n"
                for s in sample["meta"]:
                    introduction = introduction + s +'\n'
                
                for turn in sample["dialogue"]:
                    dialog_pairs = []
                    dialog_pairs.append(SPEAKER1 + f"{turn[0]}")
                    prefix += SPEAKER1 + f"{turn[0]}" +"\n"
                    if turn[1] == "":
                        dialog_pairs.append(SPEAKER2 + f"")
                        prefix += SPEAKER2 + f"" 
                    else:
                        dialog_pairs.append(SPEAKER2 + f"{turn[1]}")
                        prefix += SPEAKER2 + f"{turn[1]}" +"\n"*2
                    
            introduction = introduction + '\n'*5
            prefix = introduction + prefix
            prefix += "\n\n"

            info['prompt'] = prefix
            dialogue.append(dialog_pairs)


    #更新
    config["dialogue"] = dialogue
    # config["dialogue"] = config["dialogue"][-max_number_turns:]
    # config["user_memory"] = config["user_memory"][-max_number_turns:]


## OpenAI
# text = ''
# while text != 'quit':
#     text = input(" [Human]:")
#     if (text=='quit' or text==''):
#       break
#     prompt = prompt + ' [Human]:' + text + '\n' + ' [Robot]:' # prompt = prompt + Human+'\n'
#     print('='*72)
#     print(prompt)
#     print('='*72)
#     response = openai.Completion.create(
#       engine="text-davinci-002",
#       prompt=prompt,
#       #temperature=0.9,
#       max_tokens=150,
#       top_p=1,
#       #frequency_penalty=0.0,
#       #presence_penalty=0.6,
#       frequency_penalty=0.5,
#       presence_penalty=0.0,
#       best_of=1,
#       stop=[" [Human]:", " [Robot]:"]
#       )
#     out=response['choices'][0]["text"]
#     print(" [Topic]:",(out.lstrip('\n')))
#     prompt=prompt+out.lstrip('\n')+'\n'